{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import backtrader as bt\n",
    "import datetime\n",
    "import pandas as pd\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymysql\n",
    "\n",
    "\n",
    "# Connect to the database\n",
    "connection = pymysql.connect(host='localhost',\n",
    "                             port=3306,\n",
    "                             user='root',\n",
    "                             password='seiRaefoe9jeufooT1uipei5gungiFah',\n",
    "                             db='stock_etf',\n",
    "                             charset='utf8mb4',\n",
    "                             cursorclass=pymysql.cursors.SSDictCursor)\n",
    "\n",
    "\n",
    "try:\n",
    "    with connection.cursor() as cursor:\n",
    "        # Read a single record\n",
    "        sql = \"SELECT * FROM `etf159920_15m` WHERE datetime BETWEEN %s AND %s\"\n",
    "        cursor.execute(sql, (datetime.date(2019,1,1),datetime.date(2020,5,1)) )\n",
    "        result = cursor.fetchall()\n",
    "finally:\n",
    "    connection.close()\n",
    "\n",
    "df = pd.DataFrame(data=result)\n",
    "df = df.set_index(['datetime'])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiTFStrategy(bt.Strategy):\n",
    "    params = (\n",
    "        ('period', 20),\n",
    "    )\n",
    "    \n",
    "    # states defination\n",
    "    Empty, M15Hold, H1Hold, D1Hold = range(4)\n",
    "    States = [\n",
    "        'Empty', 'M15Hold', 'H1Hold', 'D1Hold',\n",
    "    ]\n",
    "    \n",
    "    def log(self, txt):\n",
    "        ''' Logging function for this strategy'''\n",
    "        dt = self.datas[0].datetime.datetime(0)\n",
    "        print('%s, %s' % (dt.isoformat(), txt))\n",
    "        \n",
    "    def __init__(self):\n",
    "        self.ma15m = bt.talib.SMA(self.dnames.hs15m, timeperiod=self.p.period)\n",
    "        self.ma1h = bt.talib.SMA(self.dnames.hs1h, timeperiod=self.p.period)\n",
    "        self.ma1d = bt.talib.SMA(self.dnames.hs1d, timeperiod=self.p.period)\n",
    "    \n",
    "        self.c15m = bt.indicators.CrossOver(self.dnames.hs15m, self.ma15m, plot=False)\n",
    "        self.c1h = bt.indicators.CrossOver(self.dnames.hs1h, self.ma1h, plot=False)\n",
    "        self.c1d = bt.indicators.CrossOver(self.dnames.hs1d, self.ma1d, plot=False)\n",
    "        \n",
    "        self.bsig15m = self.c15m==1\n",
    "        self.bsig1h = self.c1h==1\n",
    "        self.bsig1d = self.c1d==1\n",
    "        self.sell_signal = self.c1d==-1\n",
    "        \n",
    "        \n",
    "        self.st = self.Empty\n",
    "        self.st_map = {\n",
    "            self.Empty : self._empty,\n",
    "            self.M15Hold : self._m15hold,\n",
    "            self.H1Hold : self._h1hold,\n",
    "            self.D1Hold : self._d1hold,\n",
    "        }\n",
    "        \n",
    "        # To keep track of pending orders\n",
    "        self.order = None\n",
    "        \n",
    "    def notify_order(self, order):\n",
    "        if order.status in [order.Submitted, order.Accepted]:\n",
    "            # Buy/Sell order submitted/accepted to/by broker - Nothing to do\n",
    "            return\n",
    "\n",
    "        # Check if an order has been completed\n",
    "        # Attention: broker could reject order if not enough cash\n",
    "        if order.status == order.Completed:\n",
    "            if order.isbuy():\n",
    "                self.log(\n",
    "                    'BUY EXECUTED, St: %s, Size: %d, Price: %.2f, Cost: %.2f, Comm %.2f' %\n",
    "                    (\n",
    "                        self.States[self.st],\n",
    "                        order.executed.size,\n",
    "                        order.executed.price,\n",
    "                        order.executed.value,\n",
    "                        order.executed.comm,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            else:  # Sell\n",
    "                self.log(\n",
    "                    'SELL EXECUTED, St: %s, Size: %d, Price: %.2f, Cost: %.2f, Comm %.2f' %\n",
    "                    (\n",
    "                        self.States[self.st],\n",
    "                        order.executed.size,\n",
    "                        order.executed.price,\n",
    "                        order.executed.value,\n",
    "                        order.executed.comm\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        elif order.status in [order.Canceled, order.Margin, order.Rejected]:\n",
    "            self.log('Order Canceled/Margin/Rejected')\n",
    "\n",
    "        # Write down: no pending order\n",
    "        self.order = None\n",
    "    \n",
    "    def notify_trade(self, trade):\n",
    "        if not trade.isclosed:\n",
    "            return\n",
    "\n",
    "        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %\n",
    "                 (trade.pnl, trade.pnlcomm))\n",
    "\n",
    "    def next(self):\n",
    "        # Check if an order is pending ... if yes, we cannot send a 2nd one\n",
    "        if self.order:\n",
    "            return\n",
    "        \n",
    "        # just call state_map function\n",
    "        self.order = self.st_map[self.st]()\n",
    "        \n",
    "        # Check if we are in the market and no buy order issued\n",
    "        if self.position and not self.order:\n",
    "            # Already in the market ... we might sell\n",
    "            if self.sell_signal:\n",
    "                self.st = self.Empty\n",
    "                # Keep track of the created order to avoid a 2nd order\n",
    "                self.order = self.close()\n",
    "                \n",
    "    def _empty(self):\n",
    "        if self.bsig15m:\n",
    "            price = self.data0.close[0]\n",
    "            cash = self.broker.get_cash()\n",
    "            # 20% of the cash\n",
    "            share = int(math.floor((0.2*cash)/price))\n",
    "\n",
    "            # set state\n",
    "            self.st = self.M15Hold\n",
    "            return self.buy(size=share)\n",
    "        \n",
    "    def _m15hold(self):\n",
    "        if self.bsig1h:\n",
    "            price = self.data0.close[0]\n",
    "            cash = self.broker.get_cash()\n",
    "            # half of the remain cash ( 60% )\n",
    "            share = int(math.floor((0.5*cash)/price))\n",
    "            \n",
    "            # set state\n",
    "            self.st = self.H1Hold\n",
    "            return self.buy(size=share)\n",
    "        \n",
    "    def _h1hold(self):\n",
    "        if self.bsig1d:\n",
    "            price = self.data0.close[0]\n",
    "            cash = self.broker.get_cash()\n",
    "            # half of the remain cash (80%)\n",
    "            share = int(math.floor((0.5*cash)/price))\n",
    "            \n",
    "            # set state\n",
    "            self.st = self.D1Hold\n",
    "            return self.buy(size=share)\n",
    "        \n",
    "    def _d1hold(self):\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cerebro = bt.Cerebro(oldtrades=True)\n",
    "\n",
    "feed = bt.feeds.PandasData(dataname=df, openinterest=None, compression=15, timeframe=bt.TimeFrame.Minutes)\n",
    "\n",
    "cerebro.adddata(feed, name='hs15m')\n",
    "cerebro.resampledata(feed, name='hs1h', timeframe=bt.TimeFrame.Minutes, compression=60)\n",
    "cerebro.resampledata(feed, name='hs1d', timeframe=bt.TimeFrame.Days)\n",
    "\n",
    "cerebro.addstrategy(MultiTFStrategy)\n",
    "\n",
    "# 小场面1万起始资金\n",
    "cerebro.broker.setcash(10000.0)\n",
    "\n",
    "# 手续费万5\n",
    "cerebro.broker.setcommission(0.0005)\n",
    "\n",
    "print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())\n",
    "\n",
    "result = cerebro.run()\n",
    "\n",
    "print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cerebro.plot(\n",
    "    iplot=False,\n",
    "    start=datetime.date(2019, 11, 1),\n",
    "    end=datetime.date(2020, 1, 1),\n",
    "    style='bar',\n",
    "    barup='red',\n",
    "    bardown='green',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
